#	Euler's method solver for differential equations
#	<eyvind.niklasson@gmail.com> / http://eyvindniklasson.se
#	libraries_needed:
#	matplotlib
#	numpy
#	scipy

from math import *
import matplotlib.pyplot as plt
import os

class problem():

	def __init__(self):
		#Configuration values for the actual differential equation
		self.freq = 1 # k/m - spring constant divided by mass

	def y_primprim(self, x):
		f = -self.freq * x
		return f

	def y_prim(self, x, deriv):
		f = x + deriv
		return f

class solver():

	def __init__(self, step_size, start_coords, target_x, problem):
		#Config stuffs
		self.step_size = step_size
		self.start_coords = start_coords
		self.target_x = target_x
		self.handles = {}
		self.handles_read = {}
		self.problem = problem
		self.debug = False
		
	def load_handle_write(self,file_name):
		if not file_name in self.handles:
			if(os.path.exists(file_name)):
				os.remove(file_name)
			#Open file/files for saving
			self.handles[file_name] = open(file_name, 'wb')
		return self.handles[file_name]
		
	def load_handle_read(self,file_name):
		if not file_name in self.handles_read:
			#Open file/files for saving
			self.handles_read[file_name] = open(file_name, 'rb')
		return self.handles_read[file_name]
	
	def load(self, file_name, debug):
		if(debug and self.debug):
			print "Loading values from:", file_name
		values = self.load_handle_read(file_name + '.txt').readlines()
		self.load_handle_read(file_name + '.txt').seek(0)
		for i, value in enumerate(values):
			#print float(value.rstrip('\n'))
			values[i] = float(value.rstrip('\n'))
		return values
		
	def save(self, file_name, value, debug):
		if(debug and self.debug):
			print "Saving to file [\"" + debug + "\"]:", value
		self.load_handle_write(file_name + '.txt').write(str(value) + "\n")
	
	def euler(self):
		current_x = self.start_coords[0]
		current_y = self.start_coords[1]
		current_y_prim = self.start_coords[2]
		
		next_x = 0
		next_y = 0
		next_y_prim = 0
		
		#No of steps to take!
		no_of_steps = int(self.target_x/self.step_size + 2) # +1 for range, +1 for getting the division right :)
		print "\nRunning " + str(no_of_steps) + " no. of steps!"
		for i in range(no_of_steps):
		
			if i < 1:
			
				#First step doesn't need to be calculated
				pass
			
			if i > 0:
				next_x = current_x + self.step_size #float( format( current_x + self.step_size , '.2f') )
				next_y_prim = self.problem.y_prim(current_y_prim, self.step_size * self.problem.y_primprim(current_y))
				next_y = (current_y + self.step_size * next_y_prim)
		
				self.save('t', current_x, 'current time')
				self.save('y', current_y, 'current length')
				self.save('y_prim', current_y_prim, 'current speed')
				self.save('y_primprim', self.problem.y_primprim(current_y), 'current_acceleration')
		
				current_x = next_x
				current_y = next_y
				current_y_prim = next_y_prim
				
		for i, handle in self.handles.items():
			handle.close()

	def plot_show(self):
		#Let's show the graph as well
		plt.show()
			
	def plot(*args):
		#Initialize a new graph
		plt.figure()
		
		#Need at least x and y values to begin!
		if(len(args) < 3):
			print "Error - please provide at least two variables to plot!"
			sys.exit()
		
		#Manually define self since we are running with a variable no of variables!
		self = args[0]
		x_val = self.load(args[1], True)
		
		for i, arg in enumerate(args):
			if i > 1:
				values = self.load(arg, True)
				
				#Debug info
				if(self.debug):
					for i, x in enumerate(x_val):
						print "Plotting: ", x_val[i], values[i] 
				
				#Plot each value-set
				plt.plot(x_val, values)
		
	def custom_vars(*args):
		self = args[0]
		
		if(len(args) < 4):
			sys.exit('Error')
		
		# We are given variable names, load them from the files
		pass_args = []
		
		for arg in args[3:]:
			pass_args.append(self.load(arg, True))
		
		for var_no, var in enumerate(self.load(args[3], True)):
			
			# We only want to give one value per "variable" on in each iteration
			temp_pass_args = []
			
			for arg_no, arg in enumerate(pass_args):
				temp_pass_args.append(pass_args[arg_no][var_no])
				
			# Run the custom function
			self.save(args[1], args[2](*temp_pass_args), 'Adding custom value')
			
		for i, handle in self.handles.items():
			handle.close()
		
#Solve this using two euler approximations at the same time. Predict first derivative then second using the first. 
#Initial values - x = 0, x' = 1, x'' = 1

# Usage of custom functions:
# 1) Define custom function with variables you want from previous data
# 2) Run custom_vars(), with arg1 -> name of variable you want to save it to, arg2 -> name of custom function,
#		arg3 -> the "x" value you want to iterate with, can be any value, arg(3+n) -> any other variables you  
#		want to use in your custom function

largest_sin_error = 0
largest_energy_error = 0

def sin_compare(t, y):
	global largest_sin_error
	error = float(sin(t)) - float(y)
	if(error>largest_sin_error):
		largest_sin_error = error
	return error

def energy_level(y, y_prim):
	global largest_energy_error
	error = (0.5*y*y + 0.5*y_prim*y_prim) - 0.5
	if(error>largest_energy_error):
		largest_energy_error = error
	return error
	
#Config
step_size = 0.02													#Accuracy - "stepsize" - adjust for more/less precise plots, experiment to trade-off time/cpu-hours
start_coords = [0.0, 0.0, 1.0]										#Coordinates to start on
target_x = 10.0														#X value on which to stop

#Run it
problem = problem() 												#Define our problem

#Run it for several stepsizes
for i in range(0,4):
	print "\nDoing stepsize:", float(0.2 * pow(float(10),float(-i)))
	solver_instance = solver(float(0.2 * pow(float(10),float(-i))), start_coords, target_x, problem)		#Config our solver
	solver_instance.euler()															#Solve it!

	solver_instance.custom_vars('sin_values', sin_compare, 't', 'y')

	solver_instance.custom_vars('energy_level', energy_level, 'y', 'y_prim')

	#Plotting
	solver_instance.plot('t', 'sin_values', 'energy_level')							#Plot it!
	
	#Error calculation/checking
	largest_energy_error = 0
	print "\nFor stepsize: " + str(0.2 * pow(float(10),float(-i))) + " - largest error in energy was: "  + str(largest_energy_error) + " and largest error sin was: " + str(largest_sin_error)
	largest_sin_error = 0
	
#Should be circle if w = 1
solver_instance.plot('y', 'y_prim')							#Plot it!

#Close all open handles after we've shown the graphs!
try:
	solver_instance.plot_show()
except:
	print "Quitting!"
	for i, handle in solver_instance.handles.items():
		handle.close()
	for i, handle in solver_instance.handles_read.items():
		handle.close()